# coding:utf-8
# Author : hiicy redldw
# Date : 2019/02/27
import numpy as np
import tokenFile


def softmax(x):
    x = np.array(x)
    return np.exp(x) / np.sum(np.exp(x))


def sigmoid(z):  # 应该是预测函数
    return 1 / (1 + np.exp(-z))


def tanh(x):
    return (np.exp(x) - np.exp(-x)) / (np.exp(x) + np.exp(-x))


class Otherlstm(object):
    def __init__(self, data_dim, hidden_dim=100):
        # data_dim: 词向量维度，即词典长度; hidden_dim: 隐单元维度
        self.data_dim = data_dim
        self.hidden_dim = hidden_dim

        # 初始化各个权重
        self.whi, self.wxi, self.bi = self._init_weight()
        self.whf, self.wxf, self.bf = self._init_weight()
        self.wha, self.wxa, self.ba = self._init_weight()
        self.who, self.wxo, self.bo = self._init_weight()
        # FAQ:初始化如何设置？
        self.wy, self.by = np.random.uniform(-np.sqrt(1 / self.hidden_dim), np.sqrt(1 / self.hidden_dim),
                                             (self.data_dim, self.hidden_dim)), \
                           np.random.uniform(-np.sqrt(1 / self.hidden_dim), np.sqrt(1 / self.hidden_dim),
                                             (self.data_dim, 1))

    def _init_weight(self):
        wh = np.random.uniform(-np.sqrt(1 / self.hidden_dim), np.sqrt(1 / self.hidden_dim),
                               (self.hidden_dim, self.hidden_dim))
        wx = np.random.uniform(-np.sqrt(1 / self.data_dim), np.sqrt(1 / self.data_dim),
                               (self.hidden_dim, self.data_dim))  # 行加行 ，隐藏层对加隐藏层维度
        b = np.random.uniform(-np.sqrt(1 / self.data_dim), np.sqrt(1 / self.data_dim),
                              (self.hidden_dim, 1))  # ps:主要是为了行加行
        return wh, wx, b

    # 初始化各个状态向量,因为第一次没有h-1所以只能自己初始化
    def _init_s(self, T):
        iss = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # input gate
        fss = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # forget gate
        ass = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # current input gate
        css = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # cell state
        oss = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # output gate
        hss = np.array([np.zeros((self.hidden_dim, 1))] * (T + 1))  # hidden gate
        ys = np.array([np.zeros((self.hidden_dim, 1))] * T)  # output value
        return {'iss': iss, 'fss': fss, 'ass': ass,
                'css': css, 'oss': oss, 'hss': hss, 'ys': ys}

    # 前向传播，单个x
    def forward(self, X):
        # 向量时间长度
        T = len(X)
        ss = self._init_s(T)
        for t in range(T):  # 这里的t对应 T:时间步，T个时序状态
            # 前一个h-1输出
            ht_pre = np.array(ss['hss'][t - 1]).reshape(-1, 1)

            ss['fss'][t] = self._cal_gate(self.whf, self.wxf, self.bf, ht_pre, X[t], sigmoid)
            ss['iss'][t] = self._cal_gate(self.whi, self.wxi, self.bi, ht_pre, X[t], sigmoid)
            ss['ass'][t] = self._cal_gate(self.wha, self.wxa, self.ba, ht_pre, X[t], tanh)
            ss['oss'][t] = self._cal_gate(self.who, self.wxo, self.bo, ht_pre, X[t], sigmoid)
            ss['css'][t] = ss['fss'][t] * ss['css'][t - 1] + ss['iss'] * ss['ass'][t]
            ss['hss'][t] = ss['oss'][t] * tanh(ss['css'][t])
            ss['ys'][t] = softmax(self.wy.dot(ss['hss'][t] + self.by))
        return ss

    # 计算各个门的值
    def _cal_gate(self, wh, wx, b, h_pre, x, activate):
        return activate(wh.dot(h_pre) + wx[:, x].reshape(-1, 1) + b)

    # 预测输出，单个x
    def predict(self, x):
        ss = self.forward(x)
        # 把每个数据维度的所有单元得值展平，得预测概率最大的那个
        pre_y = np.argmax(ss['ys'].reshape(len(x), -1), axis=1)
        return pre_y

    # 计算损失函数， softmax交叉熵损失函数， (x,y)为多个样本
    def loss(self, x, y):
        cost = 0
        for i in range(len(y)):
            ss = self.forward(x[i])
            # FAQ:[range(len(y[i])), y[i]]
            pre_yi = ss['ys'][range(len(y[i])), y[i]]
            cost -= np.sum(np.log(pre_yi))
        # 统计所有y中词的个数 计算平均损失
        N = np.sum([len(yi) for yi in y])
        ave_loss = cost / N
        return ave_loss

    # 初始化偏导数
    def _init_wh_wx_grad(self, x):
        dwh = np.zeros(self.whi.shape)
        dwx = np.zeros(self.wxi.shape)
        db = np.zeros(self.bi.shape)
        return dwh, dwx, db

    #  求梯度, (x,y)为一个样本
    def bptt(self, x, y):
        dwhi, dwxi, dbi = self._init_wh_wx_grad()
        dwhf, dwxf, dbf = self._init_wh_wx_grad()
        dwho, dwxo, dbo = self._init_wh_wx_grad()
        dwha, dwxa, dba = self._init_wh_wx_grad()
        dwy, dby = np.zeros(self.wy.shape), \
                   np.zeros(self.by.shape)
        # 初始化delta_ct,因为后向传播过程，该值需要累加
        delta_ct = np.zeros((self.hidden_dim, 1))
        # 前向计算
        ss = self.forward(x)
        # 目标函数对输出y的偏导数
        delta_o = ss['ys']
        delta_o[np.arange(len(y)), y] -= 1  # FIXME:softmax求导
        for t in np.arange(len(y))[::-1]:
            # 输出层wy, by的偏导数，由于所有时刻的输出共享输出权值矩阵，故所有时刻累加
            dwy += delta_o[t].dot(ss['hss'][t].reshape(1, -1))
            dby += delta_o[t]

            # 目标函数对隐藏状态的偏导数
            # FAQ:不是说要四个状态加加吗？
            delta_ht = self.wy.T.dot(delta_o[t])
            # 各个门及状态单元的偏导数
            delta_ot = delta_ht * tanh(ss['css'][t])
            delta_ct += delta_ht * ss['oss'][t] * (1 - tanh(ss['css'][t]) ** 2)
            delta_it = delta_ct * ss['ass'][t]
            delta_ft = delta_ct * ss['css'][t - 1]
            delta_at = delta_ct * ss['iss'][t]

            delta_at_net = delta_at * (1 - ss['ass'][t] ** 2)
            delta_it_net = delta_it * ss['ass'][t] * (1 - ss['ass'][t])
            delta_ft_net = delta_ft * ss['fss'][t] * (1 - ss['fss'][t])
            delta_ot_net = delta_ft * ss['oss'][t] * (1 - ss['oss'][t])

            # 更新各权重矩阵的偏导数，由于所有时刻共享权值，故所有时刻累加
            dwhi, dwxi, dbi = self._cal_grad_delta(dwhi, dwxi, dbi, delta_it_net, ss['hss'][t - 1], x[t])
            dwhf, dwxf, dbf = self._cal_grad_delta(dwhf, dwxf, dbf, delta_ft_net, ss['hss'][t - 1], x[t])
            dwha, dwxa, dba = self._cal_grad_delta(dwha, dwxa, dba, delta_at_net, ss['hss'][t - 1], x[t])
            dwho, dwxo, dbo = self._cal_grad_delta(dwho, dwxo, dbo, delta_ot_net, ss['hss'][t - 1], x[t])

        return [dwhi, dwxi, dbi,
                dwhf, dwxf, dbf,
                dwha, dwxa, dba,
                dwho, dwxo, dbo,
                dwy, dby]

    def _cal_grad_delta(self, dwh, dwx, db, delta_net, ht_pre, x):
        dwh += delta_net * ht_pre
        dwx += delta_net * x
        db += delta_net

        return dwh, dwx, db

    # 计算梯度, (x,y)为一个样本
    def sgd_step(self, x, y, learning_rate):
        dwhf, dwxf, dbf, \
        dwhi, dwxi, dbi, \
        dwha, dwxa, dba, \
        dwho, dwxo, dbo, \
        dwy, dby = self.bptt(x, y)
        # 更新权重矩阵
        self.whf, self.wxf, self.bf = self._update_wh_wx(learning_rate, self.whf, self.wxf, self.bf, dwhf, dwxf, dbf)
        self.whi, self.wxi, self.bi = self._update_wh_wx(learning_rate, self.whi, self.wxi, self.bi, dwhi, dwxi, dbi)
        self.wha, self.wxa, self.ba = self._update_wh_wx(learning_rate, self.wha, self.wxa, self.ba, dwha, dwxa, dba)
        self.who, self.wxo, self.bo = self._update_wh_wx(learning_rate, self.who, self.wxo, self.bo, dwho, dwxo, dbo)

    def _update_wh_wx(self, lr, wh, wx, b, dwh, dwx, db):
        wh -= lr * dwh
        wx -= lr * dwx
        b -= lr * db

        return wh, wx, b

    def train(self, X_train, y_train, learning_rate=5 * 1e-3, n_epoch=5):
        losses = []
        num_examples = 0

        for epoch in range(n_epoch):
            for i in range(len(y_train)):
                self.sgd_step(X_train[i], y_train[i], learning_rate)
                num_examples += 1
            loss = self.loss(X_train, y_train)
            losses.append(loss)
            print(f'epoch{epoch+1}:loss= {loss}')
            if len(losses) > 1 and losses[-1] > losses[-2]:
                learning_rate *= 0.5
                print('decrease learning_rate to', learning_rate)


filepath = r'F:\Resources\Dataset\results-20170508-103637.csv'
dict_size = 8000
myTokenFile = tokenFile.tokenFile2vector(filepath, dict_size)
X_train, y_train, dict_words, index_of_words = myTokenFile.get_vector()
print(X_train.shape, y_train.shape, dict_words, index_of_words)
lstm = Otherlstm(dict_size, hidden_dim=100)
lstm.train(X_train[:200], y_train[:200],
           learning_rate=0.005,
           n_epoch=3)
